package mr

import "fmt"
import "log"
import "net/rpc"
import "hash/fnv"
import "time"
import "sort"
import "os"
import "errors"
import "encoding/json"
import "io/ioutil"
import "path/filepath"

//
// Map functions return a slice of KeyValue.
//
type KeyValue struct {
	Key   string
	Value string
}

// for sorting by key.
type ByKey []KeyValue

// for sorting by key.
func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

//
// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
//
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

//
// main/mrworker.go calls this function.
//
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	// Your worker implementation here

	// uncomment to send the Example RPC to the master.
	// CallExample()

	for {
		t := AskForTask(TASK_APPLY, 0)
		
		if t.WorkType == WORK_MAP {
			mapWorker(t.MapIdx, t.NReduce, t.FileName, mapf)
		}

		if t.WorkType == WORK_REDUCE {
			reduceWorker(t.ReduceIdx, t.NMap, reducef)
		}

		if t.WorkType == WORK_FINISH {
			// fmt.Println("WORK_FINISH", t)
			return
		}

		// time.Sleep(time.Second)
		time.Sleep(time.Microsecond * 100)
	}
}

// periodly ping server
func pingMaster() {
	for {
		AskForTask(PING_SERVER, 0)
		time.Sleep(time.Second)
	}
}

// read files
func mapWorker(mapIdx int, nReduce int, filename string, mapf func(string, string) []KeyValue) {
	intermediate := []KeyValue{}
	file, err := os.Open(filename)
	// log.Println("mapWorker", filename)
	if err != nil {
		log.Fatalf("cannot open %v", filename)
	}
	content, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatalf("cannot read %v", filename)
	}
	file.Close()
	kva := mapf(filename, string(content))
	intermediate = append(intermediate, kva...)

	// read data
	data := map[int][]KeyValue{}
	for i := 0; i < len(intermediate); i++ {
		key := intermediate[i].Key
		ridx := ihash(key) % nReduce
		if _, ok := data[ridx]; !ok {
			data[ridx] = []KeyValue{}
		}
		data[ridx] = append(data[ridx], intermediate[i])
	}

	// write files
	for ridx := range data {
		WriteJson(mapIdx, ridx, data[ridx])
	}

	AskForTask(MAP_FINISH, mapIdx)
}


func WriteJson(mapIdx int, reduceIdx int, data []KeyValue) bool {
	// fn := fmt.Sprintf("mr-tmp-%v-%v", mapIdx, reduceIdx)
	rn := fmt.Sprintf("mr-%v-%v", mapIdx, reduceIdx)
	fp, _ := ioutil.TempFile("", "tmp")
	enc := json.NewEncoder(fp)
	for _, kv := range data {
		err := enc.Encode(&kv)
		if err != nil {
			break
		}
	}
	
	fp.Close()
	os.Rename(fp.Name(), rn)
	// if reduceIdx == 0 {
	// 	fmt.Println("map", rn, len(data))
	// }
	return true
}

func ReadJson(fi string) ([]KeyValue, error) {
	file, err := os.Open(fi)
	defer file.Close()
	var kva = []KeyValue{}
	if err != nil {
		log.Fatalf("cannot open %v", fi)
		return nil, errors.New("Failed to open Json file.")
	}
	dec := json.NewDecoder(file)
	for {
		var kv KeyValue
		if err := dec.Decode(&kv); err != nil {
			// fmt.Printf("Error: %t %T\n", err, err)
			break
		}
		kva = append(kva, kv)
	}
	return kva, nil
}

func Exists(name string) bool {
    if _, err := os.Stat(name); err != nil {
        if os.IsNotExist(err) {
            return false
        }
    }
    return true
}

func reduceWorker(reduceIdx int, mapNum int, reducef func(string, []string) string) {
	// log.Println("reduceWorker", reduceIdx)
	pt := fmt.Sprintf("mr-[0-9]*-%v", reduceIdx) // do not match mr-out-
	files, _ := filepath.Glob(pt)
	if len(files) != mapNum {
		// fmt.Println("files != mapNum", files, mapNum, len(files))
	}
	for i := 0; i < mapNum; i++ {
		fn := fmt.Sprintf("mr-%v-%v", i, reduceIdx)
		if !Exists(fn) {
			// fmt.Println("file not exists", fn)
		}
	}

	intermediate := []KeyValue{}
	for _, fi := range files {
		res, err := ReadJson(fi)
		if err != nil { // failed
			continue
		}
		// if reduceIdx == 0 {
		// 	fmt.Println("reduce", reduceIdx, len(res), fi)
		// }
		intermediate = append(intermediate, res...)
	}

	// fmt.Println("reduce", reduceIdx, len(intermediate), files)

	sort.Sort(ByKey(intermediate))

	oname := fmt.Sprintf("mr-out-%v", reduceIdx)
	ofile, _ := os.Create(oname)

	// merge by key to reduce
	var i = 0
	for i < len(intermediate) {
		j := i + 1
		for j < len(intermediate) && intermediate[j].Key == intermediate[i].Key {
			j++
		}
		values := []string{}
		for k := i; k < j; k++ {
			values = append(values, intermediate[k].Value)
		}
		// fmt.Println(intermediate[i].Key, values)
		output := reducef(intermediate[i].Key, values)

		// this is the correct format for each line of Reduce output.
		fmt.Fprintf(ofile, "%v %v\n", intermediate[i].Key, output)

		i = j
	}

	AskForTask(REDUCE_FINISH, reduceIdx)
}

func AskForTask(cmd CommandType, taskid int) TaskReply {
	args := TaskArgs{}
	args.MsgType = cmd
	args.TaskIdx = taskid
	reply := TaskReply{}

	call("Master.HandleTask", &args, &reply)

	return reply
}

//
// example function to show how to make an RPC call to the master.
//
// the RPC argument and reply types are defined in rpc.go.
//
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	call("Master.Example", &args, &reply)

	// reply.Y should be 100.
	fmt.Printf("reply.Y %v\n", reply.Y)
}

//
// send an RPC request to the master, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := masterSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
